{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# Copyright 2019 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# =============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Jasper inference using TensorRT Inference Server\n",
    "This Jupyter notebook provides scripts to deploy high-performance inference in NVIDIA TensorRT Inference Server offering different options for the model backend, among others NVIDIA TensorRT. \n",
    "Jasper is a neural acoustic model for speech recognition. Its network architecture is designed to facilitate fast GPU inference. \n",
    "NVIDIA TensorRT Inference Server provides a datacenter and cloud inferencing solution optimized for NVIDIA GPUs. The server provides an inference service via an HTTP or gRPC endpoint, allowing remote clients to request inferencing for any number of GPU or CPU models being managed by the server\n",
    "NVIDIA TensorRT is a platform for high-performance deep learning inference. It includes a deep learning inference optimizer and runtime that delivers low latency and high-throughput for deep learning inference applications.\n",
    "## 1. Overview\n",
    "\n",
    "The Jasper model is an end-to-end neural acoustic model for automatic speech recognition (ASR) that provides near state-of-the-art results on LibriSpeech among end-to-end ASR models without any external data. The Jasper architecture of convolutional layers was designed to facilitate fast GPU inference, by allowing whole sub-blocks to be fused into a single GPU kernel. This is important for meeting strict real-time requirements of ASR systems in deployment.The results of the acoustic model are combined with the results of external language models to get the top-ranked word sequences corresponding to a given audio segment. This post-processing step is called decoding.\n",
    "\n",
    "The original paper is Jasper: An End-to-End Convolutional Neural Acoustic Model https://arxiv.org/pdf/1904.03288.pdf.\n",
    "\n",
    "### 1.1 Model architecture\n",
    "By default the model configuration is Jasper 10x5 with dense residuals. A Jasper BxR model has B blocks, each consisting of R repeating sub-blocks.\n",
    "Each sub-block applies the following operations in sequence: 1D-Convolution, Batch Normalization, ReLU activation, and Dropout. \n",
    "In the original paper Jasper is trained with masked convolutions, which masks out the padded part of an input sequence in a batch before the 1D-Convolution.\n",
    "For inference masking is not used. The reason for this is that in inference, the original mask operation does not achieve better accuracy than without the mask operation on the test and development dataset. However, no masking achieves better inference performance especially after TensorRT optimization.\n",
    "More information on the model architecture can be found in the [Jasper Pytorch README](https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechRecognition/Jasper)\n",
    "\n",
    "### 1.2 TensorRT Inference Server Overview\n",
    "\n",
    "A typical TensorRT Inference Server pipeline can be broken down into the following 8 steps:\n",
    "1. Client serializes the inference request into a message and sends it to the server (Client Send)\n",
    "2. Message travels over the network from the client to the server (Network)\n",
    "3. Message arrives at server, and is deserialized (Server Receive)\n",
    "4. Request is placed on the queue (Server Queue)\n",
    "5. Request is removed from the queue and computed (Server Compute)\n",
    "6. Completed request is serialized in a message and sent back to the client (Server Send)\n",
    "7. Completed message travels over network from the server to the client (Network)\n",
    "8. Completed message is deserialized by the client and processed as a completed inference request (Client Receive)\n",
    "\n",
    "Generally, for local clients, steps 1-4 and 6-8 will only occupy a small fraction of time, compared to steps 5-6. As backend deep learning systems like Jasper are rarely exposed directly to end users, but instead only interfacing with local front-end servers, for the sake of Jasper, we can consider that all clients are local.\n",
    "In this section, we will go over how to launch TensorRT Inference Server and client and get the best performant solution that fits your specific application needs.\n",
    "\n",
    "Note: The following instructions are run from outside the container and call `docker run` commands as required.\n",
    "\n",
    "### 1.3 Inference Pipeline in TensorRT Inference Server\n",
    "The Jasper model pipeline consists of 3 components, where each part can be customized to be a different backend: \n",
    "\n",
    "**Data preprocessor**\n",
    "\n",
    "The data processor transforms an input raw audio file into a spectrogram. By default the pipeline uses mel filter banks as spectrogram features. This part does not have any learnable weights.\n",
    "\n",
    "**Acoustic model**\n",
    "\n",
    "The acoustic model takes in the spectrogram and outputs a probability over a list of characters. This part is the most compute intensive, taking more than 90% of the entire end-to-end pipeline. The acoustic model is the only component with learnable parameters and what differentiates Jasper from other end-to-end neural speech recognition models. In the original paper, the acoustic model contains a masking operation for training (More details in [Jasper PyTorch README](https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/SpeechRecognition/Jasper/README.md)). We do not use masking for inference . \n",
    "\n",
    "**Greedy decoder**\n",
    "\n",
    "The decoder takes the probabilities over the list of characters and outputs the final transcription. Greedy decoding is a fast and simple way of doing this by always choosing the character with the maximum probability. \n",
    "\n",
    "To run a model with TensorRT, we first construct the model in PyTorch, which is then exported into a ONNX static graph. Finally, a TensorRT engine is constructed from the ONNX file and can be launched to do inference. The following table shows which backends are supported for each part along the model pipeline.\n",
    "\n",
    "|Backend\\Pipeline component|Data preprocessor|Acoustic Model|Decoder|\n",
    "|---|---|---|---|\n",
    "|PyTorch JIT|x|x|x|\n",
    "|ONNX|-|x|-|\n",
    "|TensorRT|-|x|-|\n",
    "\n",
    "In order to run inference with TensorRT outside of the inference server, refer to the [Jasper TensorRT README](https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/SpeechRecognition/Jasper/trt/README.md)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Learning objectives\n",
    "\n",
    "This notebook demonstrates:\n",
    "- Speed up Jasper Inference with TensorRT in TensorRT Inference Server\n",
    "- Use of Mixed Precision for Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Requirements\n",
    "\n",
    "Please refer to Jasper TensorRT Inference Server README.md"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Jasper Inference\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1  Prepare Working Directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "if not 'workbookDir' in globals():\n",
    "    workbookDir = os.getcwd() + \"/../\"\n",
    "print('workbookDir: ' + workbookDir)\n",
    "os.chdir(workbookDir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2  Generate TRTIS Model Checkpoints\n",
    "Use the PyTorch model checkpoint to generate all 3 model backends. You can find a pretrained checkpoint at https://ngc.nvidia.com/catalog/models/nvidia:jasperpyt_fp16.\n",
    "\n",
    "Set the following parameters:\n",
    "\n",
    "* `ARCH`: hardware architecture. use 70 for Volta, 75 for Turing.\n",
    "* `CHECKPOINT_DIR`: absolute path to model checkpoint directory.\n",
    "* `CHECKPOINT`: model checkpoint name. (default: jasper10x5dr.pt)\n",
    "* `PRECISION`: model precision. Default is using mixed precision.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env ARCH=70\n",
    "# replace with absolute path to checkpoint directory, which should include CHECKPOINT file\n",
    "%env CHECKPOINT_DIR=<CHECKPOINT_DIR> \n",
    "# CHECKPOINT file name\n",
    "%env CHECKPOINT=jasper_fp16.pt \n",
    "%env PRECISION=fp16\n",
    "!echo \"ARCH=${ARCH} CHECKPOINT_DIR=${CHECKPOINT_DIR} CHECKPOINT=${CHECKPOINT} PRECISION=${PRECISION} trtis/scripts/export_model.sh\"\n",
    "!ARCH=${ARCH} CHECKPOINT_DIR=${CHECKPOINT_DIR} CHECKPOINT=${CHECKPOINT} PRECISION=${PRECISION} trtis/scripts/export_model.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!bash trtis/scripts/prepare_model_repository.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3  Start the TensorRT Inference Server using Docker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!bash trtis/scripts/run_server.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.4. Start inference prediction in TRTIS\n",
    "\n",
    "Use the following script to run inference with TensorRT Inference Server.\n",
    "You will need to set the parameters such as: \n",
    "\n",
    "\n",
    "* `MODEL_TYPE`: Model pipeline type. Choose from [pyt, onnx, trt] for Pytorch JIT, ONNX, or TensorRT model pipeline.\n",
    "* `DATA_DIR`: absolute path to directory with audio files\n",
    "* `FILE`: relative path of audio file to `DATA_DIR`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_TYPE=\"trt\"\n",
    "DATA_DIR=os.path.join(workbookDir, \"notebooks/\")\n",
    "FILE=\"example1.wav\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!bash trtis/scripts/run_client.sh $MODEL_TYPE $DATA_DIR $FILE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can play with other examples from the 'notebooks' directory. You can also add your own audio files and generate the output text files in this way."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.5. Stop your container in the end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker stop jasper-trtis"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
